"""
Unified hook manager implementation for parameter-efficient fine-tuning (PEFT).

This module provides a comprehensive hook management framework for neural network models,
enabling parameter-efficient fine-tuning techniques such as LoRA, Prefix Tuning,
Adapter-based methods, and other PEFT approaches.

Key features:
- Flexible module detection with multiple discovery strategies
- Factory pattern for hook creation with configurable behavior
- Robust error handling and detailed logging
- Registry-based design for easy extension
- Transparent lifecycle management of hooks (creation, attachment, detachment)

The design follows SOLID principles, emphasizing modularity and extensibility.
New hook types can be registered and used without modifying existing code.
"""

import traceback
import weakref
from contextlib import contextmanager
from typing import Any, Dict, Optional, TypeVar, Union

import torch.nn as nn

from continuallearning.interfaces import HookFactoryInterface, TaskAwareHookInterface
from continuallearning.models.pefts.common import HookError
from continuallearning.models.pefts.common.config.hook_manager_config import (
    HookFactoryConfig,
    HookManagerConfig,
)
from continuallearning.models.pefts.common.utils import (
    ModuleDetector,
    create_hook_factory,
)
from continuallearning.models.pefts.hook_managers._base_hook_manager import (
    AbstractHookManager,
)
from continuallearning.registry import HOOK_MANAGER_REGISTRY
from continuallearning.utils.logging import HookSummaryReporter, get_logger

# Get logger from the new logging system
logger = get_logger(__name__)

# Type variable for module finders
T = TypeVar("T", bound=nn.Module)
H = TypeVar("H", bound=TaskAwareHookInterface)


@HOOK_MANAGER_REGISTRY.register("unified")
class UnifiedHookManager(AbstractHookManager):
    """
    Unified hook manager with flexible module detection and hook creation.

    Args:
        model: The model to which hooks will be applied
        config: Configuration for the hook manager (can be dict or HookManagerConfig)
    """

    def __init__(
        self,
        model: nn.Module,
        config: Optional[Union[Dict[str, Any], HookManagerConfig]] = None,
    ) -> None:
        super().__init__()
        # Initialize with flags
        self._initialized = False
        self._is_attached = False
        # Store model reference
        self._model_ref = weakref.ref(model)

        # Process configuration
        self.config = self._process_config(config)

        # Initialize module detector
        self.module_detector = ModuleDetector(model, self.config.discovery)

        # Initialize hook factory from config
        self.hook_factory = self._initialize_hook_factory()

        # Initialize summary reporter using the new logging system
        self.summary_reporter = HookSummaryReporter()

        # Initialize hook records
        self._hooks: Dict[str, TaskAwareHookInterface] = {}

    def _process_config(
        self, config: Optional[Union[Dict[str, Any], HookManagerConfig]]
    ) -> HookManagerConfig:
        """Process and validate configuration."""
        if config is None:
            # Use default configuration
            return HookManagerConfig()
        elif isinstance(config, HookManagerConfig):
            # Already a config object, validate it
            config.validate()
            return config
        elif isinstance(config, dict):
            # Convert dict to HookManagerConfig
            try:
                hook_manager_config = HookManagerConfig(**config)
                hook_manager_config.validate()
                return hook_manager_config
            except Exception as e:
                raise ValueError(f"Invalid configuration dictionary: {e}") from e
        else:
            raise TypeError(
                f"config must be None, dict, or HookManagerConfig, "
                f"got {type(config).__name__}"
            )

    @property
    def model(self) -> nn.Module:
        """
        Get the model instance.

        Returns:
            nn.Module: The model instance if still alive, None otherwise
        """
        if self._model_ref is None:
            raise RuntimeError("Model reference has not been set")

        model = self._model_ref()
        if model is None:
            raise RuntimeError("Model has been garbage collected")
        return model

    def _initialize_hook_factory(
        self,
    ) -> HookFactoryInterface:
        """
        Initialize hook factory from configuration.
        """
        factory_config = self.config.hook_factory

        # If it's already a HookFactoryInterface, use it directly
        if isinstance(factory_config, HookFactoryInterface):
            return factory_config

        # If it's a HookFactoryConfig, extract parameters
        if isinstance(factory_config, HookFactoryConfig):
            return create_hook_factory(
                hook_name=factory_config.hook_name,
                router_name=factory_config.router_name,
                combiner_name=factory_config.combiner_name,
                hook_config=factory_config.hook_config,
                router_config=factory_config.router_config,
                combiner_config=factory_config.combiner_config,
            )

        raise ValueError(
            f"Invalid hook_factory configuration: {type(factory_config).__name__}"
        )

    # Add logging utility method, centralizing log recording, adhering to DRY principle
    def _log_hook_operation_result(
        self,
        hook_type: str,
        module_type: str,
        path: str,
        success: bool,
        error: Optional[str] = None,
        log_level: str = "info",
    ) -> None:
        """
        Unified logging of hook operation results to log and summary reporter.

        Args:
            hook_type: Hook type name
            module_type: Module type name
            path: Module path
            success: Whether operation was successful
            error: Error message (if any)
            log_level: Log level, defaults to "info"
        """
        # Record to summary reporter
        self.summary_reporter.record_hook_creation(
            hook_type=hook_type,
            module_type=module_type,
            path=path,
            success=success,
            error=error,
        )

        # Choose appropriate log message and level based on result
        # Unified log message format
        message_template = (
            "{result} {hook_type} on '{path}' [{module_type}]{error_info}"
        )

        if success:
            # Success message
            message = message_template.format(
                result="✓",
                hook_type=hook_type,
                path=path,
                module_type=module_type,
                error_info="",
            )
            # Log at specified log level
            logger.debug(message)
        else:
            # Failure message
            error_info = f": {error}" if error else " without specific error"
            message = message_template.format(
                result="✗",
                hook_type=hook_type,
                path=path,
                module_type=module_type,
                error_info=error_info,
            )
            logger.error(message)

    @contextmanager
    def _hook_operation_context(
        self, operation: str, path: str, module_type: str, hook_type: str = "Unknown"
    ):
        """
        Provides a context manager for hook operations, unifying log and exception handling.

        This context manager provides a unified mechanism for error handling and logging,
        ensuring all hook operations receive consistent error handling.

        Args:
            operation: Operation name (e.g., "attach", "detach")
            path: Module path
            module_type: Module type name
            hook_type: Hook type name (optional)

        Yields:
            No return value, used as a context manager
        """
        try:
            logger.debug(
                f"Attempting to {operation} {hook_type} hook to '{path}' [{module_type}]"
            )
            yield
            # Operation successful, log success
            self._log_hook_operation_result(
                hook_type=hook_type,
                module_type=module_type,
                path=path,
                success=True,
                log_level="debug",
            )
        except HookError as e:
            # Handle hook-specific errors
            error_msg = str(e)
            logger.error(
                f"Hook error in {operation} for '{path}' [{module_type}]: {error_msg}"
            )
            self._log_hook_operation_result(
                hook_type=hook_type,
                module_type=module_type,
                path=path,
                success=False,
                error=error_msg,
            )
            raise
        except Exception as e:
            # Handle general exceptions
            error_msg = str(e)
            logger.error(
                f"Unexpected error in {operation} for '{path}' [{module_type}]: {error_msg}"
            )
            logger.debug(f"Detailed error: {traceback.format_exc()}")
            self._log_hook_operation_result(
                hook_type=hook_type,
                module_type=module_type,
                path=path,
                success=False,
                error=error_msg,
            )
            raise

    def add_hooks(self) -> None:
        """
        Add hooks to target modules based on configuration.

        This method performs the following steps:
        1. Initializes resources if not already done
        2. Finds target modules using the current configuration
        3. Creates hooks for each target module using the hook factory
        4. Attaches each hook and tracks it
        5. Generates a summary report of hook operations

        If no modules are found or no hooks can be attached successfully,
        appropriate warnings are logged.
        """
        if not self._initialized:
            self._initialize_resources()

        # Find target modules based on configuration
        logger.info(
            f"Searching for target modules in model {type(self.model).__name__}"
        )
        target_modules = self.module_detector.find_target_modules()

        if not target_modules:
            logger.warning("No target modules found for hooking")
            return

        logger.info(
            f"Found {len(target_modules)} potential target modules, applying hooks..."
        )

        # Reset summary reporter
        self.summary_reporter.reset()

        # Create and attach hooks for each target module
        success_count = 0
        for idx, module_dict in enumerate(target_modules):
            path, module = module_dict["path"], module_dict["module"]
            module_type = type(module).__name__
            module_id = id(module)

            # Skip if this module already has a hook
            if path in self._hooks:
                logger.debug(
                    f"Module '{path}' [{module_type}] already hooked, skipping"
                )
                continue

            # Changed from INFO to DEBUG level to reduce log volume
            logger.debug(
                f"[{idx + 1}/{len(target_modules)}] Creating and attaching hook for '{path}' [{module_type}, id={module_id}]"
            )

            try:
                # Create hook and check type compatibility
                hook = self.hook_factory(path, module)
                hook_type = type(hook).__name__

                logger.debug(
                    f"Successfully created {hook_type} hook instance, checking compatibility..."
                )

                # Verify the hook implements TaskAwareHookInterface
                if not isinstance(hook, TaskAwareHookInterface):
                    error_msg = f"Hook factory produced hook type {hook_type} that doesn't implement TaskAwareHookInterface"
                    logger.error(error_msg)
                    self._log_hook_operation_result(
                        hook_type=hook_type,
                        module_type=module_type,
                        path=path,
                        success=False,
                        error=error_msg,
                    )
                    raise TypeError(error_msg)

                # Use context manager to handle hook attachment process
                with self._hook_operation_context(
                    "attach", path, module_type, hook_type
                ):
                    _ = hook.attach()
                    self._hooks[path] = hook
                    success_count += 1

                # Log attachment success - Changed to DEBUG level, keeping only one concise success log per module
                logger.debug(
                    f"✓ Successfully attached {hook_type} to '{path}' [{module_type}]"
                )
                self._log_hook_operation_result(
                    hook_type=hook_type,
                    module_type=module_type,
                    path=path,
                    success=True,
                )

            except (HookError, TypeError, Exception):
                # Exceptions already handled in context manager, no need to re-log here
                pass

        if success_count > 0:
            # Generate and display summary report - Keep summary reports at INFO level
            summary = self.summary_reporter.get_summary_report()
            logger.info(summary)

            self._is_attached = True
        else:
            logger.warning(
                "No hooks were successfully attached, check configuration or model structure"
            )

    def remove_hooks(self) -> None:
        """
        Remove all hooks from the model.

        This method safely detaches all hooks managed by this hook manager,
        cleaning up resources and restoring the model to its original state.
        It tracks success/failure of each detachment operation and generates
        a summary report.
        """
        if not self._hooks:
            logger.debug("No hooks to remove")
            return

        # Reset summary reporter
        self.summary_reporter.reset()

        # Record total hook count for accurate reporting after clearing
        total_hooks = len(self._hooks)
        logger.info(f"Removing all hooks, total: {total_hooks}")

        success_count = 0
        for idx, (path, hook) in enumerate(list(self._hooks.items())):
            module_type = type(hook.module).__name__
            hook_type = type(hook).__name__

            # Use context manager to handle hook detachment process
            try:
                with self._hook_operation_context(
                    "detach", path, module_type, hook_type
                ):
                    hook.detach()
                    success_count += 1
            except Exception:
                # Exceptions already handled in context manager
                pass

        self._hooks.clear()
        self._is_attached = False
        logger.info(
            f"Removed all hooks, {success_count}/{total_hooks} successfully detached"
        )
